1852B - Imbalanced Arrays - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;

int b[100010];


void linearSort(deque<int>& ids, vector<int>& a) {
    int n = ids.size();
    vector<int> count(*max_element(a.begin(), a.end()) + 1, 0);
    vector<int> output(n);

    // Count the occurrences of each value in 'a'
    for (int i = 0; i < n; i++) {
        count[a[ids[i]]]++;
    }

    // Calculate the cumulative count
    for (int i = 1; i < count.size(); i++) {
        count[i] += count[i - 1];
    }

    // Build the output vector
    for (int i = n - 1; i >= 0; i--) {
        output[count[a[ids[i]]] - 1] = ids[i];
        count[a[ids[i]]]--;
    }

    // Copy the sorted IDs back to the 'ids' vector
    for (int i = 0; i < n; i++) {
        ids[i] = output[i];
    }
}





int main(){
    ios::sync_with_stdio(0);
    cin.tie(nullptr);
    int Ncase;
    cin>>Ncase;
    while(Ncase--){
        int n;
        cin>>n;
        vector<int> a(n,0);
        deque<int> ids;
        for(int i=0;i<n;i++){
            int temp;
            cin>>temp;
            a[i]=temp;
            ids.push_back(i);
        }
        // sort(ids.begin(), ids.end(), [a](int l, int r){
        //     return a[l]<a[r];
        // });
        linearSort(ids,a);
        int del_cnt=0,p=n;
        while(!ids.empty() && (((a[ids.front()]-del_cnt)==0)^((a[ids.back()]-del_cnt)==ids.size()))){
            if((a[ids.front()]-del_cnt)==0){
                b[ids.front()] = -p;
                ids.pop_front();
            }
            else{
                b[ids.back()] = p;
                ids.pop_back();
                del_cnt++;
            }
            p--;
        }
        if(!ids.empty()){
            cout<<"NO"<<endl;
            continue;
        }

        cout<<"YES"<<endl;
        for(int i=0;i<n;i++){
            cout<<b[i]<<" ";
        }
        cout<<endl;

    }
}